
import os

from langchain_google_genai import GoogleGenerativeAIEmbeddings

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "playground"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["GOOGLE_API_KEY"] = "AIzaSyBJoz7BvdFgWTBwzcu-0xWpJKfEJOR6vPM"

embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")

def simple_emb_demo():

    vector = embeddings.embed_query("hello, world!")
    print(len(vector))
    print(vector[:5])


def simple_batch_demo():
    vectors = embeddings.embed_documents(
        [
            "Today is Monday",
            "Today is Tuesday",
            "Today is April Fools day",
        ]
    )
    print(vectors)
    print(len(vectors[0]))


"""
    task_type 有以下几种类型：
    task_type_unspecified
    retrieval_query
    retrieval_document
    semantic_similarity
    classification
    clustering
"""
def simple_emb_demo3():
    query ='你好'
    query_2 = '我好'
    answer_1 = '他好'
    query_embeddings = GoogleGenerativeAIEmbeddings(
        model="models/embedding-001", task_type="retrieval_query"
    )
    doc_embeddings = GoogleGenerativeAIEmbeddings(
        model="models/embedding-001", task_type="retrieval_document"
    )
    query_vecs = [query_embeddings.embed_query(q) for q in [query, query_2, answer_1]]
    print(query_vecs)
    doc_vecs = [doc_embeddings.embed_query(q) for q in [query, query_2, answer_1]]
    print(doc_vecs)


if __name__ == "__main__":
    simple_emb_demo3()